/*
 * ibdtrace.cpp
 *
 *  Created on: Sep 20, 2011
 *      Author: lindenb
 */
#include <fstream>
#include <cerrno>
#include <cmath>
#include <cassert>
#include <iomanip>
#include "xstdlib.h"
#include "ibdtrace.h"

using namespace std;
using namespace bdb;



ApplicationTracer::ApplicationTracer():PS(NULL)
	{
	pair_height=100;
	margin_top=100;
	margin_bottom=100;
	margin_right=100;
	margin_left=100;
	x_axis_width=1000;
	genome_size=0;
	pair_margin_y=5;
	}

ApplicationTracer::~ApplicationTracer()
	{
	if(PS!=NULL)  delete PS;
	}

void ApplicationTracer::buildChromInfo()
	{
	genome_size=0L;
	//cerr<< __FUNCTION__ << ":" << __LINE__ << endl;
	auto_ptr<Cursor> c=env()->linkageDB->cursor();
	while(c->next())
		{
		Position pos;
		c->keyCopyTo(&pos,0,sizeof(Position));

		if(chromInfos.empty() || chromInfos.back().chrom!=pos.chrom)
			{
			//cerr<< __FUNCTION__ << ":" << __LINE__ << ":" <<(int) pos.chrom << endl;
			ChromInfo info;
			info.chrom=pos.chrom;
			info.minPos=pos.pos;
			info.maxPos=pos.pos;
			info.fraction=0.0f;
			info.width=1.0f;

			chromInfos.push_back(info);
			}
		else
			{
			ChromInfo& info=chromInfos.back();

			info.minPos=std::min(pos.pos,info.minPos);
			info.maxPos=std::max(pos.pos,info.maxPos);
			}
		}
	c.reset();
	//cerr<< __FUNCTION__ << ":" << __LINE__ << endl;
	finalizeChromInfo();
	}


void ApplicationTracer::finalizeChromInfo()
	{
	genome_size=0L;
	//cerr<< __FUNCTION__ << ":" << __LINE__ << endl;
	for(size_t i=0;i< chromInfos.size();++i)
		{
		ChromInfo& info=chromInfos[i];
		position_t shift=(position_t)((info.length()/100.0)*5.0);
		info.minPos=max(info.minPos-shift,0);
		info.maxPos=info.maxPos+shift;
		genome_size+=(int64_t)info.length();
		//cerr << "info " << (int)info.chrom<<":" << info.minPos << "-"<< info.maxPos << endl;
		}
	//cerr<< __FUNCTION__ << ":" << __LINE__ << endl;
	pixel_t x=0;
	for(size_t i=0;i< chromInfos.size();++i)
		{
		ChromInfo& info=chromInfos[i];
		position_t len=((position_t)info.maxPos-(position_t)info.minPos);
		info.fraction=(len/(double)genome_size);
		info.width= x_axis_width*info.fraction;
		info.x=x;
		x+=info.width;
		}
	//cerr<< __FUNCTION__ << ":" << __LINE__ << endl;
	}


void ApplicationTracer::buildIndividual(const Individual* indi)
	{
	const int pvalue_index=2;
	if(PS==NULL)
		{
		PS=new Postscript(*output);
		PS->out() << "%!PS-Adobe-3.0"
				<<endl;
		}

	/* collect all pair with this individual */
	vector<PairInfo> pairInfos;
	individual_id_t indi_id=indi->id;
	set<pair_id_t> observed_pairs;
	auto_ptr<Cursor> c=env()->pairOfIndividualsToPairId->cursor();
	while(c->next())
		{
		PairOfIndividual pairIndi;

		c->keyCopyTo(&pairIndi,0,sizeof(PairOfIndividual));

		if(!pairIndi.contains(indi_id))
			{
			continue;
			}

 		pair_id_t pair_id;
		c->dataCopyTo(&pair_id,0,sizeof(pair_id_t));
		if(observed_pairs.insert(pair_id).second)//inserted
			{
			PairInfo pi;
			pi.other=pairIndi.complementary(indi_id);
			pi.index=pairInfos.size();
			pi.pair_id=pair_id;
			pi.data=NULL;
			pi.flag_is_digest=false;
			pairInfos.push_back(pi);
			}
		}
	c.reset();


	/* TODO: sort pairInfos */

	/* add extra track*/

	PairInfo* pairInfoDigest=NULL;
	if(true)//TODO
		{
		PairInfo pi;
		pi.other=-1;
		pi.index=pairInfos.size();
		pi.pair_id=-1;
		pi.data=NULL;
		pi.flag_is_digest=true;
		pairInfos.push_back(pi);
		pairInfoDigest=&pairInfos.back();
		}

	/* loop over the pairs , set the y axis*/
	pixel_t y=0;
	for(size_t i=0;i<pairInfos.size();++i)
		{
		if(i>0) y+=pair_margin_y;
		pairInfos.at(i).y=y;
		y+=pair_height;
		}


	const pixel_t y_axis_height=y;


	PS->out() << "%%DocumentMedia: Custom "
					<< (int)(margin_left+margin_right+x_axis_width) << " "//width
					<<(int)(y_axis_height+margin_top+margin_bottom) << " " //height
					<< " 0 () ()\n"
				"%%Title: IBD\n"
				"%%Creator: "<< __FILE__ << endl <<
				"%%Pages: (atend)\n"
				"%%Orientation: Landscape\n"
				<< "%%BoundingBox: 0 0 "
					<< (int)(margin_left+margin_right+x_axis_width)<< " "
					<< (int)(y_axis_height+margin_top+margin_bottom)<< endl<<
				"%%EndComments\n"
					;
	PS->out() << "/Segment { /points exch def points 0 get points 1 get moveto "
					"/index 2 def points 2 2 points length 1 sub{ points index get "
					"/index index 1 add def "
					"points index get "
					"lineto } for }  def\n";



	Genotype* genotypeRow=NULL;

	/* print main sample */
	{
	ostringstream os;
	os << indi->family->name;
	os << " : ";
	os << indi->name;
	string sampleName(os.str());
	PS->moveto(10,margin_bottom+y_axis_height+20).
		setfont("Courier",14).
		setgray(0.1).
		show(sampleName.c_str());
	}

	/* print sample label */
	for(vector<PairInfo>::iterator ri=pairInfos.begin();
					 ri!=pairInfos.end();++ri)
			{
			ostringstream os;
			if(ri->flag_is_digest)
				{
				os << "*";
				}
			else
				{
				Individual* other= env()->id2individual[ri->other];
				os << other->family->name;
				os << " ";
				os << other->name;
				}
			string sampleName(os.str());

			PS->moveto(
					1,
					margin_bottom+ri->y+pair_height/2.0).
				setgray(0.1).
				setfont("Courier",10).
				show(sampleName.c_str());
			}


    /* loop over each chromosome */
	for(size_t cidx=0;cidx< chromInfos.size();++cidx)
		{
		map<pair_id_t,vector<PositionXY>* > pair2positionXY;
		ChromInfo& chromInfo=chromInfos[cidx];

		GState g(this->PS);
		if(chromInfo.width>200)
			{
			PS->setgray(0.1);
			for(int i=0;i < 10;++i)
				{
				position_t tick= (position_t)(chromInfo.minPos+(chromInfo.length()/10.0)*i);
				pixel_t x=margin_left + (chromInfo.width/10.0)*i;
				char tmp[50];
				snprintf(tmp,50,"%d",(int)tick);
				PS->line(x,margin_bottom,x,margin_bottom-5);
				PS->stroke();
				PS->moveto(x,margin_bottom-10);
				PS->rotate(-90);
				PS->show(tmp);
				PS->rotate(90);
				}
			}



		PS->translate(margin_left+chromInfo.x,margin_bottom);
		PS->newpath().
			rect(0,0,chromInfo.width,y_axis_height).
			setgray(cidx%2==0?0.9:0.8).
			closepath().
			fill();

		string chromName=env()->id2chromosome[chromInfos[cidx].chrom];
		PS->moveto((chromInfo.width-(chromName.size()*12))/2.0,y_axis_height+1).
			setfont("Courier",12).
			setgray(0.1).
			show(chromName.c_str());

		Position position;
		position.chrom=chromInfo.chrom;
		position.pos=chromInfo.minPos;

		auto_ptr<Cursor> c=ibdEnv.linkageDB->cursor(NULL,(u_int32_t)0);
		bool first=true;
		/** loop over each marker on this chromosome */
		for(;;)
			{
			if(first)
				{
				first=false;
				DBT dbtK;
				DBT dbtV;
				memset(&dbtK,0,sizeof(DBT));
				memset(&dbtV,0,sizeof(DBT));
				dbtK.data=(void*)&position;
				dbtK.size=sizeof(Position);
				if(c->get(&dbtK,&dbtV,DB_SET_RANGE)!=0)
					{
					break;
					}
				}
			else
				{
				if(!c->next())
					{
					break;
					}
				}
			Position curr;
			c->keyCopyTo((void*)&curr,0,sizeof(Position));

			if(curr.chrom!=chromInfo.chrom ||
			   curr.pos> chromInfo.maxPos)
				{
				//cerr << "break" << curr << " vs " << chromInfo.minPos<<"-"<< chromInfo.maxPos<< endl;
				break;
				}
			if(genotypeRow==NULL) genotypeRow=(Genotype*)safeMalloc(env()->count_pairs*sizeof(Genotype));
			c->dataCopyTo((void*)genotypeRow,0,env()->count_pairs*sizeof(Genotype));



			for(vector<PairInfo>::iterator ri=pairInfos.begin();
				 ri!=pairInfos.end();++ri)
				{
				if(ri->flag_is_digest) continue;
				if(ri->pair_id >= env()->count_pairs)
					{
					THROW("runtime error" << ri->pair_id << " > "<< env()->count_pairs);
					}


				if(ri->data==NULL)
					{
					ri->data=new vector<PositionXY>;
					}

				PositionXY posxy;
				posxy.position=curr.pos;
				posxy.pvalue=genotypeRow[ri->pair_id].pvalue[pvalue_index];

				assert(ri->data->empty() || ri->data->back().position<=curr.pos);
				ri->data->push_back(posxy);
				}
			}//end of scan linkage for this chrom

		if(pairInfoDigest!=NULL && pairInfos.size()>1)
			{
			pairInfoDigest->data=new vector<PositionXY>;
			size_t n_markers=0;
			for(vector<PairInfo>::iterator ri=pairInfos.begin();
							 ri!=pairInfos.end();++ri)
				{
				if(ri->flag_is_digest || ri->data==NULL) continue;
				if(n_markers==0)
					{
					n_markers=ri->data->size();
					}
				else
					{
					assert(n_markers==ri->data->size());
					}
				}


			for(size_t i=0;i< n_markers;++i)
				{
				int count=0;
				position_t pos=-1;
				pvalue_t sum=0;
				for(vector<PairInfo>::iterator ri=pairInfos.begin();
						 ri!=pairInfos.end();++ri)
					{
					if(ri->flag_is_digest || ri->data==NULL) continue;

					if(ri->data->at(i).pvalue>0.9)
						{
						if(pos==-1) pos= ri->data->at(i).position;
						sum+=ri->data->at(i).pvalue;
						count++;
						}
					}
				if(count>1 && pos!=-1)
					{
					PositionXY posxy;
					posxy.position=pos;
					posxy.pvalue=(sum/count);
					pairInfoDigest->data->push_back(posxy);
					}
				}
			}


		//digest data
		for(vector<PairInfo>::iterator ri=pairInfos.begin();
				 ri!=pairInfos.end();++ri)
			{
			if(ri->data==NULL)
				{
				continue;
				}

			vector<Postscript::Point> points;
			for(size_t n=0;n< ri->data->size();++n)
				{
				Postscript::Point pt;
				assert(ri->data->at(n).position>=chromInfo.minPos );
				assert(ri->data->at(n).position<=chromInfo.maxPos );

				pt.x= ((ri->data->at(n).position- chromInfo.minPos)/(pixel_t)(chromInfo.length()))*chromInfo.width;
				pt.y= ri->data->at(n).pvalue* (this->pair_height-1);
				if( !points.empty() &&
					fabs((float)(points.back().x-pt.x))<0.1f &&
					fabs((float)(points.back().y-pt.y))<0.1f
					)
					{
					//cerr << fabs((float)(points.front().x-pt.x)) << " " << fabs((float)(points.front().y-pt.y)) << "\n";
					//ignore
					}
				else if(points.size()>1 &&
						fabs((float)(points[points.size()-2].y-points[points.size()-1].y))< 0.01 &&
						fabs((float)(points[points.size()-2].y-pt.y))< 0.01
						)
					{
					points[points.size()-1]=pt;
					}
				else
					{

					points.push_back(pt);
					}
				}
			delete ri->data;
			ri->data=NULL;

			if(points.size()>1)
				{
				GState g2(PS);//shift vertical for pair track
				PS->translate(0,ri->y);
				PS->newpath().
					rect(0,0,chromInfo.width,pair_height).
					line(0,pair_height/2.0f,chromInfo.width,pair_height/2.0f).//IBD 1
					setlinewidth(0.1).
					closepath().
					stroke();

				if(!ri->flag_is_digest)
					{
					for(size_t n=0;n< points.size();++n)
						{
						if(n==0) PS->moveto(points[n].x,points[n].y);
						else PS->lineto(points[n].x,points[n].y);
						}

					PS->setgray(0.2).setlinewidth(0.4).stroke();
					}
				else
					{
					for(size_t n=0;n< points.size();++n)
						{
						PS->line(points[n].x,0,points[n].x,points[n].y);
						}
					PS->red().setlinewidth(0.7).stroke();
					}
				}


			}
		}/* end loop chromosome */
	free(genotypeRow);




	PS->showpage();
	PS->out() << "%%Trailer\n"

			;
	}


void ApplicationTracer::usage(std::ostream& out)
	{
	out << "options\n" <<
		" --db-home (berkeydb home) default $DB_HOME " << (db_home==NULL?"":db_home)<< " \n" <<
		" --family (family)\n"
		" --individual (family)\n"
		" --position chrom:start-end\n"
		" --width (int)\n"
		" --height (int)\n"
		" -o <output-file> (or stdout if omitted).\n"
		<< std::endl
		;
	}

int ApplicationTracer::main(int argc,char** argv)
	{
	if(argc==1)
		{
		usage(cout);
		return EXIT_SUCCESS;
		}

	int optind=1;
	char* position=NULL;
	char* fileout=NULL;
	auto_ptr<fstream> oustream;
	char* family_name=NULL;
	char* individual_name=NULL;
	while(optind < argc)
		{

		if(std::strcmp(argv[optind],"-h")==0)
		    {
		    usage(cout);
		    std::exit(EXIT_FAILURE);
		    }
		else if(std::strcmp(argv[optind],"--db-home")==0 && optind+1<argc)
		    {
		    db_home=argv[++optind];
		    }
		else if(std::strcmp(argv[optind],"--position")==0 && optind+1<argc)
		    {
		    position=argv[++optind];
		    }
		else if(std::strcmp(argv[optind],"--family")==0 && optind+1<argc)
		    {
			family_name=argv[++optind];
		    }
		else if(std::strcmp(argv[optind],"--individual")==0 && optind+1<argc)
			{
			individual_name=argv[++optind];
			}
		else if(std::strcmp(argv[optind],"--width")==0 && optind+1<argc)
			{
			x_axis_width=max(50,atoi(argv[++optind]));
			}
		else if(std::strcmp(argv[optind],"--height")==0 && optind+1<argc)
			{
			pair_height=max(50,atoi(argv[++optind]));
			}
		else if(std::strcmp(argv[optind],"-o")==0 && optind+1<argc)
			{
			fileout=argv[++optind];
			}
		else if(std::strcmp(argv[optind],"--debug")==0)
			{
			ibdEnv.debug=true;
			}
		else if(std::strcmp(argv[optind],"--")==0)
			{
			++optind;
			break;
			}
		else if(argv[optind][0]=='-')
			{
			std::cerr << "unknown option " << argv[optind] << std::endl;
			std::exit(EXIT_FAILURE);
			}
		else
			{
			break;
			}
		++optind;
		}
	if(checkHome()!=0) return EXIT_FAILURE;
	if(optind!=argc)
	    {
	    cerr << "Illegal number of arguments"<< endl;
	    return EXIT_FAILURE;
	    }
	if(family_name==NULL)
		{
		 cerr << "Undefined family name"<< endl;
		 return EXIT_FAILURE;
		}
	if(individual_name==NULL)
		{
		cerr << "Undefined individual name"<< endl;
		return EXIT_FAILURE;
		}
	if(fileout!=NULL)
		{
		oustream.reset(new fstream(fileout,ios::out));
		if(!oustream->is_open())
			{
			cerr << "Cannot open \""<< fileout << "\" " << strerror(errno) << endl;
			return EXIT_FAILURE;
			}
		this->output=oustream.get();
		}
	//cerr << "Open" << endl;
	ibdEnv.open(db_home,true);
	//cerr << "Build" << endl;
	buildChromInfo();



	if(position!=NULL)
		{
		auto_ptr<GenomicPosition> ptr= GenomicPosition::parse(position);
		if(ptr.get()==NULL)
			{
			cerr << "Cannot parse position "<< position << endl;
			return EXIT_FAILURE;
			}
		map<string,tid_t>::iterator rc=ibdEnv.chrom2tid.find(ptr->chromosome);
		if(rc==ibdEnv.chrom2tid.end())
			{
			cerr << "undefined chromosome "<< position << endl;
			return EXIT_FAILURE;
			}
		chromInfos.clear();
		ChromInfo ci;
		ci.chrom=rc->second;
		ci.fraction=100.0f;
		ci.x=0;
		ci.width=x_axis_width;
		ci.minPos=ptr->chromStart;
		ci.maxPos=ptr->chromEnd;
		chromInfos.push_back(ci);
		finalizeChromInfo();
		}


	map<string,Family*>::iterator r1=ibdEnv.name2family.find(family_name);
	if(r1==ibdEnv.name2family.end())
		{
		cerr << "Unknown family name "<< family_name<< endl;
		return EXIT_FAILURE;
		}
	Individual* indi=r1->second->findIndividualByName(individual_name);
	if(indi==NULL)
		{
		cerr << "Unknown indi name "<< individual_name<< endl;
		return EXIT_FAILURE;
		}

	buildIndividual(indi);

	ibdEnv.close();
	if(oustream.get()!=NULL)
		{
		oustream->flush();
		oustream->close();
		this->output=&cout;
		oustream.reset();
		}
	return EXIT_SUCCESS;
	}
